import json

import datasets
from datasets import DownloadManager, DatasetInfo


class CMRC2018TRIAL(datasets.GeneratorBasedBuilder):

	def _info(self) -> DatasetInfo:
		"""
				info方法, 定义数据集的信息,这里要对数据的字段进行定义
		:return:
		"""
		return datasets.DatasetInfo(
			description = "CMRC2018 trial",
			features = datasets.Features({
				"id": datasets.Value("string"),
				"context": datasets.Value("string"),
				"question": datasets.Value("string"),
				"answers": datasets.features.Sequence(
					{
						"text": datasets.Value("string"),
						"answer_start": datasets.Value("int32"),
					}
				)
			})
		)

	def _split_generators(self, dl_manager: DownloadManager):
		"""
				返回datasets.SplitGenerator
				涉及两个参数: name和gen_kwargs
				name: 指定数据集的划分
				gen_kwargs: 指定要读取的文件的路径, 与_generate_examples的入参数一致
		:param dl_manager:
		:return: [ datasets.SplitGenerator ]
		"""
		return [datasets.SplitGenerator(name = datasets.Split.TRAIN,
			gen_kwargs = {"filepath": "./cmrc2018_trial.json"})]

	def _generate_examples(self, filepath):
		"""
				生成具体的样本, 使用yield
				需要额外指定key, id从0开始自增就可以
		:param filepath:
		:return:
		"""
		# Yields (key, example) tuples from the dataset
		with open(filepath, encoding = "utf-8") as f:
			data = json.load(f)
			for example in data["data"]:
				for paragraph in example["paragraphs"]:
					context = paragraph["context"].strip()
					for qa in paragraph["qas"]:
						question = qa["question"].strip()
						id_ = qa["id"]

						answer_starts = [answer["answer_start"] for answer in qa["answers"]]
						answers = [answer["text"].strip() for answer in qa["answers"]]

						yield id_, {
							"context": context,
							"question": question,
							"id": id_,
							"answers": {
								"answer_start": answer_starts,
								"text": answers,
							},
						}